from parse.stmt import *
from parse.expr import *
from parse.resolver import Scope, ScopeType, Variable, VariableType
from parse.token import TokenType
from vm.chunk import Chunk
from vm.value import Value, ValueType, Obj, ObjFunc, ObjClosure, ObjClass
from vm.instruction import OpCode
from vm.machine import VM, CallFrame
from typing import Optional


class CodeGenerator(StmtVisitor, ExprVisitor):
    def __init__(self, stmts: List[Stmt]):
        self.stmts = stmts
        self.compiler: Optional[Compiler] = None
        self.class_compiler: Optional['ClassCompiler'] = None
        self.line = 1
        self.vm = VM()

    def generate_code(self):
        main_fun_stmt = StmtFun(Token('__main__', TokenType.FUN, 0), [], self.stmts, FunType.FUN)
        self.begin_func_scope(main_fun_stmt)
        for stmt in self.stmts:
            stmt.accept(self)
        self.emit_ins(OpCode.CONSTANT)
        self.emit_constant(Value.nil())
        self.emit_ins(OpCode.RETURN)
        obj_val = self.end_func_scope()
        fn = obj_val.value.func_value()

        call_frame = CallFrame(ObjClosure(fn, []), [])
        self.vm.frames.append(call_frame)

        return self.vm

    def visit_print_stmt(self, stmt: 'StmtPrint'):
        stmt.expr.accept(self)
        self.emit_ins(OpCode.PRINT)
        self.emit_ins(OpCode.POP)

    def visit_expr_stmt(self, stmt: 'StmtExpr'):
        stmt.expr.accept(self)
        self.emit_ins(OpCode.POP)

    def visit_var_stmt(self, stmt: 'StmtVar'):
        self.line = stmt.name.line
        self.declare_variable(stmt.name.lexeme, stmt.initializer, stmt.local)

    def visit_block_stmt(self, stmt: 'StmtBlock'):
        self.begin_scope()
        for s in stmt.stmts:
            s.accept(self)
        self.end_scope()

    def visit_if_stmt(self, stmt: 'StmtIf'):
        stmt.condition.accept(self)
        else_jump = self.emit_jump(OpCode.JUMP_IF_FALSE)
        self.emit_ins(OpCode.POP)
        stmt.then_stmt.accept(self)
        then_jump = self.emit_jump(OpCode.JUMP)
        self.compiler.patch_jump(else_jump)
        self.emit_ins(OpCode.POP)
        if stmt.else_stmt:
            stmt.else_stmt.accept(self)
        self.compiler.patch_jump(then_jump)

    def visit_while_stmt(self, stmt: 'StmtWhile'):
        start_point = self.begin_loop_scope()
        stmt.condition.accept(self)
        exit_jump = self.emit_jump(OpCode.JUMP_IF_FALSE)
        self.emit_ins(OpCode.POP)
        stmt.body.accept(self)

        # 跳转到循环体头部
        loop_jump = self.emit_jump(OpCode.JUMP)
        self.compiler.patch_jump(loop_jump, start_point)

        # 跳出循环的点
        self.compiler.patch_jump(exit_jump)
        self.emit_ins(OpCode.POP)
        self.end_loop_scope()

    def visit_break_stmt(self, stmt: 'StmtBreak'):
        loop_break_point = self.emit_jump(OpCode.JUMP)
        self.compiler.register_break_point(loop_break_point)

    def visit_continue_stmt(self):
        loop_continue_point = self.emit_jump(OpCode.JUMP)
        self.register_loop_continue_point(loop_continue_point)

    def visit_class_stmt(self, stmt: 'StmtClass'):
        # 创建一个class对象，将其压栈
        self.begin_class_scope(stmt)
        self.emit_ins(OpCode.CLASS)
        self.emit_constant(Value.string(stmt.name.lexeme))
        if stmt.superclass is not None:
            if stmt.name.lexeme == stmt.superclass.lexeme:
                raise ValueError(f"cannot inherit self")

            self.emit_get_variable(stmt.superclass.lexeme)
            self.emit_ins(OpCode.INHERIT)
            self.emit_ins(OpCode.POP)

        for s in stmt.methods:
            s.accept(self)
            self.emit_ins(OpCode.METHOD)
            self.emit_ins(OpCode.POP)

        if self.in_global_scope():
            self.emit_ins(OpCode.DEFINE_GLOBAL)
            self.emit_constant_string(stmt.name.lexeme)
            self.emit_ins(OpCode.POP)
        else:
            self.compiler.add_variable(stmt.name.lexeme, VariableType.LOCAL)

        self.end_class_scope()

    def visit_function_stmt(self, stmt: 'StmtFun'):
        """
        虽然有虚拟机和字节码，但clox的实现其实还是个解释器，因为源代码编译后就直接加载到内存中，变成可执行的映像，
        而没有可执行文件或者类似Java的.class字节码文件的中间结果
        """
        self.begin_func_scope(stmt)
        if stmt.fun_type == FunType.METHOD:
            self.compiler.add_variable('this', VariableType.LOCAL)
        for param in stmt.params:
            self.compiler.add_variable(param.lexeme, VariableType.LOCAL)
        for s in stmt.body:
            s.accept(self)

        obj_val = self.end_func_scope()

        # 该指令生成闭包对象，捕捉变量，并将闭包对象压至栈顶
        self.emit_ins(OpCode.CLOSURE)
        self.emit_constant(obj_val)
        for v in self.compiler.vars_to_nested_scopes():
            if v.var_type == VariableType.FREE or v.var_type == VariableType.UPVALUE:
                self.emit(v.var_type, v.offset)
            else:
                raise TypeError(f'invalid captured variable type: {v.var_type}')

        # 函数作为变量
        if self.in_global_scope() and stmt.fun_type != FunType.METHOD:
            self.emit_ins(OpCode.DEFINE_GLOBAL)
            self.emit_constant_string(stmt.name.lexeme)
            self.emit_ins(OpCode.POP)
        else:
            self.compiler.add_variable(stmt.name.lexeme, VariableType.LOCAL)

    def visit_return_stmt(self, stmt: 'StmtReturn'):
        stmt.expr.accept(self)
        self.emit_ins(OpCode.RETURN)

    def visit_binary(self, expr: 'ExprBinary'):
        expr.left.accept(self)
        expr.right.accept(self)

        match expr.op.token_type:
            case TokenType.PLUS:
                self.emit_ins(OpCode.ADD)
            case TokenType.MINUS:
                self.emit_ins(OpCode.SUBTRACT)
            case TokenType.STAR:
                self.emit_ins(OpCode.MULTIPLY)
            case TokenType.SLASH:
                self.emit_ins(OpCode.DIVIDE)
            case TokenType.GREATER:
                self.emit_ins(OpCode.GREATER)
            case TokenType.GREATER_EQUAL:
                self.emit_ins(OpCode.LESS)
                self.emit_ins(OpCode.NOT)
            case TokenType.LESS:
                self.emit_ins(OpCode.LESS)
            case TokenType.LESS:
                self.emit_ins(OpCode.GREATER)
                self.emit_ins(OpCode.NOT)
            case _:
                raise ValueError(f'invalid binary op {expr.op}')

    def visit_grouping(self, expr: 'ExprGrouping'):
        expr.expr.accept(self)

    def visit_literal(self, expr: 'ExprLiteral'):
        self.emit_ins(OpCode.CONSTANT)
        if isinstance(expr.value, str):
            self.emit_constant(Value(ValueType.OBJ, Obj.string(expr.value)))
        elif isinstance(expr.value, float):
            self.emit_constant(Value.new_number(expr.value))
        elif isinstance(expr.value, bool):
            self.emit_constant(Value.boolean(expr.value))
        elif expr.value is None:
            self.emit_constant(Value.nil())
        else:
            raise ValueError(f'invalid literal value {expr.value}')

    def visit_unary(self, expr: 'ExprUnary'):
        expr.expr.accept(self)
        match expr.token.token_type:
            case TokenType.BANG:
                self.emit_ins(OpCode.NOT)
            case TokenType.MINUS:
                self.emit_ins(OpCode.NEGATE)

    def visit_variable(self, expr: 'ExprVariable'):
        self.emit_get_variable(expr.name.lexeme)

    def visit_assign(self, expr: 'ExprAssign'):
        expr.expr.accept(self)
        var = self.resolve_variable(expr.name.lexeme)
        if var is None:
            self.emit_ins(OpCode.SET_GLOBAL)
            self.emit_constant_string(expr.name.lexeme)
        elif var.var_type == VariableType.LOCAL:
            self.emit_ins(OpCode.SET_LOCAL)
            self.emit(var.offset)
        elif var.var_type == VariableType.FREE:
            self.emit_ins(OpCode.SET_FREE_VALUE)
            self.emit(var.offset)
        else:
            self.emit_ins(OpCode.SET_UPVALUE)
            self.emit(var.offset)

    def visit_logical(self, expr: 'ExprLogical'):
        expr.left.accept(self)
        if expr.op.token_type == TokenType.AND:
            shortcircuit_jmp = self.emit_jump(OpCode.JUMP_IF_FALSE)
            self.emit_ins(OpCode.POP)
            expr.right.accept(self)
            self.compiler.patch_jump(shortcircuit_jmp)
        else:
            jmp2 = self.emit_jump(OpCode.JUMP_IF_FALSE)
            self.emit_ins(OpCode.POP)
            shortcircuit_jmp = self.emit_jump(OpCode.JUMP)
            self.emit_ins(OpCode.POP)
            self.compiler.patch_jump(jmp2)
            expr.right.accept(self)
            self.compiler.patch_jump(shortcircuit_jmp)

    def visit_call(self, expr: 'ExprCall'):
        #  将闭包对象压栈
        expr.callee.accept(self)

        #  将参数压栈
        for arg in expr.arguments:
            arg.accept(self)

        self.emit_ins(OpCode.CALL)
        #  参数个数
        self.emit(len(expr.arguments))

    def visit_get(self, expr: 'ExprGet'):
        expr.obj.accept(self)
        self.emit_ins(OpCode.GET_PROPERTY)
        self.emit_constant_string(expr.name.lexeme)

    def visit_set(self, expr: 'ExprSet'):
        expr.obj.accept(self)
        expr.value.accept(self)
        self.emit_ins(OpCode.SET_PROPERTY)
        self.emit_constant_string(expr.name.lexeme)

    def visit_this(self, expr: 'This'):
        var = self.resolve_variable('this')
        self.emit_ins(OpCode.GET_LOCAL)
        self.emit(var.offset)

    def visit_super(self, expr: 'Super'):
        if self.class_compiler is None:
            raise ValueError(f'cannot visit super out of a class')
        if self.class_compiler.super_class is None:
            raise ValueError(f'class {self.class_compiler.name} does not have super')
        self.emit_get_variable(self.class_compiler.super_class)
        self.emit_ins(OpCode.GET_SUPER)
        self.emit_constant_string(expr.token.lexeme)

    def emit(self, *ins: int):
        for i in ins:
            # print(f'emit {i}')
            self.compiler.chunk.write(i, self.line)

    def emit_get_variable(self, name: str):
        var = self.resolve_variable(name)
        if var is None:
            self.emit_ins(OpCode.GET_GLOBAL)
            self.emit_constant_string(name)
        elif var.var_type == VariableType.LOCAL:
            self.emit_ins(OpCode.GET_LOCAL)
            self.emit(var.offset)
        elif var.var_type == VariableType.FREE:
            self.emit_ins(OpCode.GET_FREE_VALUE)
            self.emit(var.offset)
        elif var.var_type == VariableType.UPVALUE:
            self.emit_ins(OpCode.GET_UPVALUE)
            self.emit(var.offset)

    def emit_constant(self, v: Value):
        # print(f'emit constant {v}')
        self.compiler.chunk.add_constant(v, self.line)

    def emit_jump(self, ins: OpCode):
        self.emit_ins(ins)
        self.emit(0xff, 0xff)
        return len(self.compiler.chunk.code) - 2

    def emit_constant_string(self, s: str):
        # print(f'emit string {s}')
        s_obj = self.vm.copy_string(s)
        self.compiler.chunk.add_constant(Value(ValueType.OBJ, s_obj), self.line)

    def declare_variable(self, name: str, expr: 'Expr', local: bool):
        if self.in_global_scope():
            expr.accept(self)
            self.emit_ins(OpCode.DEFINE_GLOBAL)
            self.emit_constant_string(name)
            self.emit_ins(OpCode.POP)
        else:
            expr.accept(self)
            if not local:
                self.emit_ins(OpCode.ESCAPE_VALUE)
                var_type = VariableType.FREE
            else:
                var_type = VariableType.LOCAL

            self.append_variable_to_scope(name, var_type)

    def begin_scope(self):
        self.compiler.begin_scope()

    def end_scope(self):
        self.compiler.end_scope()

    def begin_loop_scope(self):
        return self.compiler.begin_loop_scope()

    def end_loop_scope(self):
        self.compiler.end_loop_scope()

    def begin_func_scope(self, stmt: StmtFun):
        """
        将闭包捕捉到的自由变量放到当前scope中
        """
        if self.compiler:
            captured_vars = self.compiler.vars_to_nested_scopes()
            start_depth = self.compiler.current_depth + 1
        else:
            captured_vars = []
            start_depth = 0

        compiler = Compiler(stmt.name.lexeme, start_depth, captured_vars, len(stmt.params))
        compiler.enclosing = self.compiler
        self.compiler = compiler

        scope_type = ScopeType.FUNC if stmt.fun_type == FunType.FUN else FunType.METHOD

        self.compiler.begin_scope(scope_type)

    def end_func_scope(self) -> Value:
        self.compiler.end_scope()
        obj_fn = self.compiler.finish()
        self.compiler = self.compiler.enclosing

        val = Value(ValueType.OBJ, obj_fn)

        self.vm.add_obj(val.value)

        return val

    def begin_class_scope(self, stmt: StmtClass):
        if stmt.superclass is not None:
            super_class = stmt.superclass.lexeme
        else:
            super_class = None
        class_compiler = ClassCompiler(stmt.name.lexeme, self.class_compiler, super_class)
        self.class_compiler = class_compiler

    def end_class_scope(self):
        self.class_compiler = self.class_compiler.enclosing

    def register_loop_start_point(self):
        return self.compiler.register_loop_start_point()

    def register_loop_continue_point(self, loop_continue_point: int):
        self.compiler.resolve_loop_scope().loop_continue_points.append(loop_continue_point)

    def resolve_variable(self, lexeme: str) -> Variable:
        return self.compiler.resolve_variable(lexeme)

    def append_variable_to_scope(self, lexeme: str, var_type: VariableType):
        self.compiler.add_variable(lexeme, var_type)

    def emit_ins(self, ins: OpCode):
        # print(f'emit op {ins.name}')
        self.compiler.chunk.write(ins, self.line)

    def in_global_scope(self):
        return self.compiler.is_in_global


class ClassCompiler:
    def __init__(self, name: str, enclosing: Optional['ClassCompiler'], super_class: Optional[str]):
        self.name = name
        self.enclosing = enclosing
        self.super_class = super_class


class Compiler:
    """
    基本的编译单位，函数或者方法
    """
    def __init__(self,
                 name: str,
                 base_depth: int,
                 # 从外层scope捕获到的自由变量
                 inherited_vars: List[Variable],
                 arity: int):
        self._curr_depth = base_depth - 1
        self.inherited_vars = inherited_vars
        self.name = name
        self.chunk = Chunk(name)
        self.scopes: List[Scope] = []
        self.enclosing: Optional['Compiler'] = None
        self.arity = arity

    def finish(self):
        return ObjFunc(self.name, self.arity, len(self.vars_to_nested_scopes()), self.chunk)

    @property
    def base_depth(self):
        return self.scopes[-1].depth

    def begin_scope(self, scope_type: ScopeType = ScopeType.PLAIN):
        self._curr_depth += 1
        self.scopes.append(Scope(self._curr_depth, scope_type))

    def begin_loop_scope(self):
        self.begin_scope(ScopeType.LOOP)
        return self.register_loop_start_point()

    def end_loop_scope(self):
        scope = self.resolve_loop_scope()
        for break_point in scope.loop_break_points:
            self.patch_jump(break_point)

        for continue_point in scope.loop_continue_points:
            self.patch_jump(continue_point, scope.loop_start)
        self.end_scope()

    def patch_jump(self, origin_offset: int, target_pos=None):
        """
        修正在位置origin_offset处跳转指令的跳转偏移
        :param origin_offset: 跳转指令的偏移量所在位置
        :param target_pos: 跳转目标位置，默认为当前
        :return:
        """
        if target_pos is None:
            delta = len(self.chunk.code) - origin_offset - 2
        else:
            delta = target_pos - origin_offset - 2

        self.chunk.code[origin_offset] = (delta >> 8) & 0xff
        self.chunk.code[origin_offset + 1] = delta & 0xff

    def register_break_point(self, loop_break_point: int):
        loop_scope = self.resolve_loop_scope()
        if loop_scope is None:
            raise ValueError('not loop match break statement')
        loop_scope.loop_break_points.append(loop_break_point)

    @property
    def current_depth(self):
        return self._curr_depth

    def resolve_loop_scope(self):
        for scope in reversed(self.scopes):
            if scope.scope_type == ScopeType.LOOP:
                return scope

        return None

    def register_loop_start_point(self):
        scope = self.resolve_loop_scope()
        scope.loop_start = len(self.chunk.code)
        return scope.loop_start

    def register_loop_continue_point(self, loop_continue_point: int):
        self.resolve_loop_scope().loop_continue_points.append(loop_continue_point)

    def resolve_variable(self, lexeme: str) -> Optional[Variable]:
        for scope in reversed(self.scopes):
            v = scope.resolve_variable(lexeme)
            if v:
                return v

        return self.resolve_up_value(lexeme)

    def resolve_up_value(self, name: str) -> Optional[Variable]:
        for offset, v in enumerate(self.inherited_vars):
            if v.name == name:
                # closure_vars中的offset是上层scope的，用于内部使用的offset直接使用下标
                return Variable(v.name, offset, VariableType.UPVALUE)

        return None

    @property
    def is_in_global(self):
        return self.scopes[-1].depth == 0

    def end_scope(self):
        self.scopes.pop()

    def add_variable(self, name: str, var_type: VariableType):
        """
        :param name: 变量名
        :param var_type: 变量类型，此处仅有local或者free两种
        :return:
        """
        if name not in self.scopes[-1]:
            self.scopes[-1].add_variable(name, var_type=var_type)
        else:
            raise ValueError(f'duplicated variable {name}')

    def vars_to_nested_scopes(self):
        ret = []

        for v in reversed(self.scopes):
            for v1 in v.variables:
                if v1.var_type == VariableType.FREE:
                    ret.append(v1)

        for v in self.inherited_vars:
            ret.append(v)

        return ret

    def is_func(self):
        return self.scopes[-1].scope_type == ScopeType.FUNC

    def is_method(self):
        return self.scopes[-1].scope_type == ScopeType.METHOD
